from __future__ import division
import numpy as np
#from numpy import multiply as nm
from ShallowDonor import Shallow_Donor
from Bulk import Silicon
from Constant_Physics import Physics
from Tools import *
from MRTools import *
from MRTools import Lorentzian as lorentz
import matplotlib.pyplot as plt
import time
from scipy.integrate import odeint
import src_fortran.froutines as fr
from termcolor import colored
from scipy.optimize import curve_fit

#To Do list
#    - B sweep, with good linewidth
#    - fitting parameters to a B/f sweep
#    - export (.csv/.ps/...)

class SDR_parameters:
    #SDR model parameters
    PE = 1e3
    Cp = 1e4
    Cap = 1e4 
    Rp = 2e4
    Rap = 1e6
    
    #Radiation characteristics
    mw_rate=1e10 #the microwave intensity
    mw_frequency=9076#MHz
    
    #Static magnetic field
    B0=0.066056
    B0xyz=np.array([0, 0, 1])
    
    #Modulation
    modulation_type="field" #field/AM/FM
    modulation_frequency=100000 #Hz
    modulation_phase=0 #rad
    modulation_shape="sin" #sin/square
    modulation_amplitude=0.003 #T/ percent/ MHz

    #Signal properties        
    width_MHz=5
    width_mT=1
    
    #Experimental conditions
    temperature=16#K 

##########################################
##########################################

class Bfield(Physics):
    
    def __init__(self,B0=0.1, xyz=np.array([0, 0, 1]), mod_freq=100000, mod_phase=0,mod_form="sin",mod_amp=0.001):
        self.B0=B0*Physics.FieldUnit
        self.xyz=xyz
        self.mod_freq=mod_freq
        self.mod_phase=mod_phase
        self.mod_form=mod_form
        self.mod_amp=mod_amp*Physics.FieldUnit
    
    def field_vec(self,t):
        """
        Input:
            -t: the time
        Output:
            -field vector
        """
        o=np.zeros([3,1])
        return fr.bfield(o,self.B0,self.xyz,self.mod_freq,self.mod_phase,self.mod_amp,t)
        
    def field_norm(self,t):
        """
        Input:
            -t: the time
        Output:
            -field norm
        """
        if self.mod_form=="sin":
            return (self.B0 + self.mod_amp* np.sin(2*np.pi*t*self.mod_freq+self.mod_phase))
        elif self.mod_form=="no_mod":
            return self.B0

class ResonanceFrequency(Physics):
    
    def __init__(self,f0=0.1, mod_freq=100000, mod_phase=0,mod_form="sin",mod_amp=0.001):
        """
        Input:
            -f0:
            -mod_freq:
            -mod_phase:
            -mod_form:
            -mod_amp:
        """
        self.f=f0*Physics.FieldUnit
        self.mod_freq=mod_freq
        self.mod_phase=mod_phase
        self.mod_form=mod_form
        self.mod_amp=mod_amp*Physics.FieldUnit
        
    def frequency(self,t):
        """
        Input:
            -t: the time
        Output:
            -frequency
        """
        if self.mod_form=="sin":
            return (self.f + self.mod_amp* np.sin(2*np.pi*t*self.mod_freq+self.mod_phase))
        elif self.mod_form=="no_mod":
            return self.f
        
class ResonanceFrequencyAmplitude(Physics):
    
    def __init__(self,A0=0.1, mod_freq=100000, mod_phase=0,mod_form="sin",mod_amp=1):
        """
        Input:
            -f0:
            -mod_freq:
            -mod_phase:
            -mod_form:
            -mod_amp: in percent
        """
        self.A=A0*Physics.FieldUnit
        self.mod_freq=mod_freq
        self.mod_phase=mod_phase
        self.mod_form=mod_form
        self.mod_amp=mod_amp*Physics.FieldUnit
        
    def amplitude(self,t):
        """
        Input:
            -t: the time
        Output:
            -frequency
        """
        if self.mod_form=="sin":
            return (self.A + self.mod_amp* np.sin(2*np.pi*t*self.mod_freq+self.mod_phase))
        elif self.mod_form=="no_mod":
            return self.A

##########################################
##########################################
class Lockin_Detection():
    def __init__(self,fundamental_f=100000, n_harmonics=1):
        """
        Inputs:
            -fundamental_f: in Hz
            -n_harmonics: integer, 1-> 2*f0, 2-> 3*f0
        """
        self.f0=fundamental_f
        self.harmonics=np.arange(1,n_harmonics+2,1)      
        self.signals_X={}
        self.signals_Y={}  
        self.signals_R={}
        self.signals_phi={}  
        print self.harmonics
        
    def detect_XY(self,time_ax,signal):
        """
        Input:
            -time_ax: the corresponding time axis, np.array([n])
            -signal: np.array([n])
        """
        for harm in self.harmonics:
            print harm
            f_det=harm*self.f0
            phi0 = np.sum(signal*np.sin(2*np.pi*f_det*time_ax))
            phi90 = np.sum(signal*np.sin(2*np.pi*f_det*time_ax+np.pi/2))
            self.signals_X[harm]=phi0
            self.signals_Y[harm]=phi90

    def calculate_Rphi(self):
        for harm in self.harmonics:
            self.signals_R[harm]=np.sqrt(self.signals_X[harm]**2+self.signals_Y[harm]**2)
            # arctan2(y,x)
            self.signals_phi[harm]=np.arctan2(self.signals_Y[harm],self.signals_X[harm])


##########################################
##########################################    
class Transient(SDR_parameters):

    def __init__(self,sdr_param):
        pass

    def is_steady_state_reached(self):
        """
        Description:
            Print whether the steady state is reached or not
        Todo:
            -Redefine the steady state
        Input:
            -odeint_sol: solution of the system
            -time_scale=t.size
        Output:
            -None
        Modification of properties:
            -None
        Dependence:
            -Properties:
                -self.conductivity
                -self.steadystate_margin
                -self.is_steady_state_reached_print
        """
        self.CT.TIC("is_steady_state_reached")
        if self.conductivity==[]:
            print "\nPlease compute the conductivity first."
        else:
            if np.abs(self.conductivity[-1]-self.conductivity[-2])<=self.steadystate_margin:
                if self.is_steady_state_reached_print==True:
                    print "\nSteady state reached: ds/dt= ",np.abs(self.conductivity[-1]-self.conductivity[-2])
            else:
                if self.is_steady_state_reached_print==True:
                    print "\nSteady state not reached yet: ds/dt= ",np.abs(self.conductivity[-1]-self.conductivity[-2])
                self.is_steady_state_reached_flag=False
        self.CT.TOC("is_steady_state_reached")

    def Plot_transient(self,transient_type="conductivity,light resonance"):
        """
        Input:
            - transient_type:    "conductivity": plot the conductivity transient
        """
        
        #Determing the number of plots
        current_subplot=1
        n_plots=1
        if "resonance" not in transient_type:
            n_plots=1
        elif "resonance" in transient_type:
            n_plots=2
            if "fft" in transient_type:
                n_plots+=1
        
        
        fig = plt.figure()
        ax=fig.add_subplot(n_plots,1,current_subplot)
        ax.plot(self.t1,self.conductivity_0_ss,'b')
        ax.set_title('Light on - transient')
        ax.set_ylabel('SDR(a.u.)', color='b')
        for tl in ax.get_yticklabels():
            tl.set_color('b')

        if "field" in transient_type:
            self.fieldrecordt1=np.zeros([self.t1.size])
            for i in range(self.t1.size):
                self.fieldrecordt1[i]=self.field.field_norm(self.t1[i])
            ax2 = ax.twinx()
            ax2.plot(self.t1, self.fieldrecordt1, 'r')
            ax2.set_ylabel('Field(T)', color='r')
            for tl in ax2.get_yticklabels():
                tl.set_color('r')
        current_subplot+=1

                
        if "resonance" in transient_type:
            ax=fig.add_subplot(n_plots,1,current_subplot)
            ax.plot(self.t2,self.conductivity_ss_mwss)
            ax.set_title('Resonance on - transient')
            ax.set_ylabel('SDR(a.u.)', color='b')
            for tl in ax.get_yticklabels():
                tl.set_color('b')
            if "field" in transient_type:
                self.fieldrecordt2=np.zeros([self.t2.size])
                for i in range(self.t2.size):
                    self.fieldrecordt2[i]=self.field.field_norm(self.t2[i])
                ax2 = ax.twinx()
                ax2.plot(self.t2, self.fieldrecordt2, 'r')
                ax2.set_ylabel('Field', color='r')
                for tl in ax2.get_yticklabels():
                    tl.set_color('r')
            current_subplot+=1
            if "fft" in transient_type:
                ax=fig.add_subplot(n_plots,1,current_subplot)
                self.timescale_with_resonance_ss=int(1/20*self.t2.size)
                ffty=np.abs(np.fft.fft(self.conductivity_ss_mwss[SDR.timescale_with_resonance_ss:-1]))
                n=ffty.size
                freq = np.fft.fftfreq(n, d=self.time_intervale_with_resonance)
                ax.plot(freq, ffty)#, freq, sp.imag)
                ax.set_title('Resonance on - transient')
                ax.set_ylabel('fft SDR(a.u.)', color='b')
            
        plt.show()
        
    def Temporal_evolution_without_resonance(self,sdr_model,starting_time_step=0):
        """
        Description:
            Resolve the time dependent problem without resonance (solving dN/dt=self.NonResonantRates*N)
        Todo:
            -None
        Input:
            -None
        Output:
            -None
        Modification of properties:
            -self.NonResonantRates
            -self.steady_conductivity
        Dependence:
            -Properties:
                -self.full_dim
                -self.c3_i,self.c3_e
                -self.dim_n
                -self.timescale_without_resonance,self.time_intervale_without_resonance
                -self.steadystate
                -self.conductivity
            -Methods:
                -self.Non_Resonant_Rates_Evaluation()
                -self.Compute_conductivity(y,t.size)
                -self.is_steady_state_reached()
        """
        self.CT.TIC("Temporal_evolution_without_resonance")
        #Print information
        print colored("Temporal evolution without any resonance:",'blue')
        
        self.previous_step_t=starting_time_step
        self.Do_Eig(time_step=starting_time_step)
        
        self.Non_Resonant_Rates_Evaluation() #Compute the rates matrix for the non resonant case
        
        # preparing the initial state, no nuclear polarization
        
        y0=np.zeros([self.full_dim])
        y0[self.c3_i:self.c3_e]=1/self.dim_n*np.ones([self.dim_n])
        
        # defining the time scale
        self.t1=Arange(self.last_time_step,self.last_time_step+self.timescale_without_resonance,self.time_intervale_without_resonance)
        # defining the function dy/dt = f(y,t)
        def func(y,t0):
            #Don't know why, but some iterations don't increase the time step
            if self.previous_step_t<t0:
                self.Do_Eig(time_step=t0)
                self.Non_Resonant_Rates_Evaluation() #Compute the rates matrix for the non resonant case

            self.previous_step_t=t0
            Print_progress(t0,self.t1[-1],start=self.t1[0])
        
            return np.dot(self.NonResonantRates,y)

        
        # performing the ODEINT
        y=odeint(func,y0,self.t1,hmax=self.time_intervale_without_resonance)
        self.last_time_step=self.t1[-1]
        
        self.conductivity_0_ss=self.Compute_conductivity(y,self.t1.size)
        self.Check_list()
  
        #taking the last solution step (last line) and save it as steadystate, and storing the corresponding conductivity
        self.steadystate = y[-1,:]
        self.steady_conductivity=self.conductivity[-1]
        self.CT.TOC("Temporal_evolution_without_resonance")

    def Temporal_evolution_with_resonance(self,sdr_model,starting_time_step=0,excitation_freq=[[9076,5]]):
        """
        Description:
            Resolve the time dependent problem at resonance (solving dN/dt=self.ResonantRates*N)
        Todo:
            -None
        Input:
            -None
        Output:
            -None
        Modification of properties:
            -self.ResonantRates
            -self.resonant_conductivity
        Dependence:
            -Properties:
                -self.timescale_with_resonance,self.time_intervale_with_resonance
                -self.steadystate
                -self.conductivity
            -Methods:
                -self.Resonant_Rates_Evaluation()
                -self.Compute_conductivity(y,t.size)
                -self.is_steady_state_reached()
        """
        self.CT.TIC("Temporal_evolution_with_resonance")
        #Print information
        print colored("\nTemporal evolution with resonance(s):",'blue')

        #Compute the rates matrix for the resonant case
        # if self.coupling_PRC_donor==False,, we compute it at each time step (resource consuming?)
        self.previous_step_t=starting_time_step
        
        self.Do_Eig(time_step=starting_time_step)
                
        self.Resonant_Rates_Evaluation(excitation_freq=excitation_freq)
        
        # preparing the initial state, no nuclear polarization
        y0=self.steadystate
        
        # defining the time scale
        self.t2=Arange(self.last_time_step,self.last_time_step+self.timescale_with_resonance,self.time_intervale_with_resonance)

        # defining the function dy/dt = f(y,t)
        def func(y,t0):
            #Don't know why, but some iterations don't increase the time step
            if self.previous_step_t<t0:
                self.Do_Eig(time_step=t0)
                #Compute the rates matrix for the resonant case
                # if self.coupling_PRC_donor==False,, we compute it at each time step (resource consuming?)
                self.Resonant_Rates_Evaluation(excitation_freq=excitation_freq)
            self.previous_step_t=t0
            Print_progress(t0,self.t2[-1],start=self.t2[0])
            return np.dot(self.ResonantRates,y)
            
        # performing the ODEINT
        y=odeint(func,y0,self.t2,hmax=self.time_intervale_with_resonance)
        self.last_time_step=self.t2[-1]
        
        self.conductivity_ss_mwss=self.Compute_conductivity(y,self.t2.size)#Only the last row conductivity
        self.Check_list()
        self.resoningstate = y[-1,:]
        self.resonant_conductivity=self.conductivity[-1]
        self.CT.TOC("Temporal_evolution_with_resonance")
        
    def fit_transient(self,t,ft):
        """
        We suppose ft to be exp
        We want to know "tau"
        """
        def func(t,A,tau):
            return A*np.exp(-t/tau)
        popt,pcov=curve_fit(func,t,ft)
        return popt
    
##########################################
##########################################
class Parametrized_SDR_model(SDR_parameters):
    
    def __init__(self, bulk, specie,clk):
        """
        specie: define tje shallow donor specie ("P", "As", "Sb5/2", "Sb7/2", "Bi")
        D: Dissociation rate
        G: Generation rate
        C: Capture rate
        Rp: Recombination rate of parallel pairs
        Rap: Recombination rate of antiparallel pairs
        W: carrier generation due to illumination
        temperature: in Kelvin
        average_dist: average separation between one donor and one PRC, in 'nm'
        """
        self.clk=clk
        self.clk.TIC("Initialization")
        self.fortran_compiled_functions=True
        
        self.steadystate_margin=1e-6
        self.is_steady_state_reached_flag=True
        self.is_steady_state_reached_print=False
        self.non_resonant_matrix_pop_conservative=True
        self.resonant_matrix_pop_conservative=True
        self.coupling_PRC_donor=True
        
        self.print_population=False

        self.donor = specie
        self.bulk = bulk
        self.PRCdim = 2
        self.PRCspin = 1 / 2
        self.electronicdim = 2
        self.nucleardim = int(2 * self.donor.nspin + 1)
        #-----------------
        self.dim_een = int(self.PRCdim * self.electronicdim * self.nucleardim) #number of states with the donor ionized
        self.dim_n = int(self.nucleardim) #number of states with both ionized
        self.full_dim = self.dim_een + self.dim_een + self.dim_n

        self.Mbj = np.zeros([self.dim_een, self.dim_een])
        self.Mjb = np.zeros([self.dim_een, self.dim_een])
        self.Mbj_el_wise_square = np.zeros([self.dim_een, self.dim_een])
        self.Mjb_el_wise_square = np.zeros([self.dim_een, self.dim_een])
        #-----------------
        
        # Rate matrix: column names
            #[0...c1_e, c2_i=c1_e...c2_e, c3_i=c2_e...c3_e]
            #[                                                ]
            #[                                                ]
            #[                                                ]
            #[                                                ]
        self.c1_i=0
        self.c1_e=self.dim_een
        self.c2_i=self.c1_e
        self.c2_e=self.dim_een + self.dim_een
        self.c3_i=self.c2_e
        self.c3_e=self.full_dim+1
        
        self.timescale_without_resonance=1e-3
        self.time_intervale_without_resonance=1e-3*self.timescale_without_resonance
        self.timescale_with_resonance=1e-4
        self.time_intervale_with_resonance=1e-4*self.timescale_with_resonance
        self.conductivity=[]
        self.steadystate=[]
        self.resoningstate=[]
        self.steady_conductivity=0
        
        self.dipolar_transitions=[]    # [[B0,li-lj,EPR,SDR],...]


        self.norm=0
        
        #This is a marker to remember the time during the simulation
        self.last_time_step=0


        # Non Resonant rate equations
        self.NonResonantRates = np.zeros([self.full_dim, self.full_dim])
        # Resonant rate equations
        self.ResonantRates = np.zeros([self.full_dim, self.full_dim])
        # Non resonant b term
        self.population_away_from_resonance = np.zeros(self.dim_een)
        # Resonant b term
        self.population_at_resonance = np.zeros(self.dim_een)


        #Construction of the eigenvectors lists and hamiltonians matrix
        self.eigvalDpRz = np.zeros([self.dim_een])
        self.eigvalDzRz = np.zeros([self.dim_een])
        self.hamiltonianDzRz = np.zeros([self.dim_een, self.dim_een])
        self.eigvalDpRm = np.zeros([self.dim_n])
        self.clk.TOC("Initialization")
        
    def Non_Resonant_Rates_Evaluation(self):
        """
        Description:
            Calculate the Non Resonant rates matrix
            20120911,9am: execution time around 65 ms
        To do:
            -include the relaxation (and its dependence on the magnetic field)
        Input:
            - t: the time
        Output:
            -None
        Modification of properties:
            -self.NonResonantRates
        Dependences:
            -Properties:
                -self.full_dim
                -self.Cp
                -self.Cap
                -self.PE: photoexcitqtion rate
                -self.Rp
                -self.Rap
                -self.dim_een
                -self.dim_n
            -Class:
                -self.donor
                -self.bulk
        """
        
        self.clk.TIC("Non_Resonant_Rates_Evaluation")
        if self.fortran_compiled_functions:
            o=np.zeros([self.dim_een+self.dim_een+self.dim_n,self.dim_een+self.dim_een+self.dim_n])
            self.NonResonantRates=fr.nonresmat(o,
                                           self.Mbj_el_wise_square,
                                           self.dim_een+self.dim_een+self.dim_n,
                                           self.dim_een,
                                           self.bulk.dim_PRC,
                                           self.donor.dim_e,
                                           self.dim_n,
                                           self.Cp,
                                           self.Cap,
                                           self.Rp,
                                           self.Rap,
                                           self.PE)
        else:
            #Construct self.NRrates
                # Ci: assuming there is no coupling at this step (even if we write the transformation matrix)
            Calpha = (self.Cp - self.Cap) / 2
            Cbeta = (self.Cp + self.Cap) / 2
            Ci = np.zeros([self.dim_een, self.dim_een])
            Ci = Calpha * np.kron(np.kron(self.bulk.Sz, self.donor.Sz), self.donor.Id) + Cbeta * np.kron(np.kron(self.bulk.Sd, self.donor.Sd), self.donor.Id)
            
                # Cj
            Cj = np.zeros([self.dim_een, self.dim_een])
            v=np.kron(np.array([self.Cp,self.Cap,self.Cap,self.Cp]),np.ones([self.dim_n]))
            for k in range(self.dim_een):
                Cj[k,:]=(v*self.Mbj_el_wise_square[:,k]).T
                #print v.size
                #print (v*self.Mbj_el_wise_square[:,k]).size
                #print Cj[k,:].size
    
                # Rj
            Rj = np.zeros([self.dim_een, self.dim_een])
            R2=np.array([self.Rp,self.Rap,self.Rap,self.Rp])
            R3=np.ones([self.dim_n])
            R1=np.kron(R2,R3)
            for j in range(self.dim_een):
                Rj[j,j]=np.dot(R1.T,self.Mbj_el_wise_square[:,j])
    
                # Rk
            Rk = np.zeros([self.dim_n, self.dim_een])
            for j in range(self.dim_n):
                R4=np.zeros([self.dim_n])
                R4[j]=1
                R4=np.kron(R2,R4)
                for i in range(self.dim_een):
                    Rk[j,i]=np.dot(R4,self.Mbj_el_wise_square[:,i])
                
            
            
                # W: we assume there is no polarization of the excitation light ...
            PEi = self.PE/self.dim_een*np.kron(np.ones([self.PRCdim*self.electronicdim,1]),self.donor.Id)
            PEk = self.PE / self.dim_n * self.donor.Id
    
                # Construct the rates matrix
            self.Ci=Ci
            self.Cj=Cj
            self.Rj=Rj
            self.Rk=Rk
            self.PEk=PEk
            self.PEi=PEi
            self.NonResonantRates[self.c1_i:self.c1_e, self.c1_i:self.c1_e] = -Ci
            self.NonResonantRates[self.c2_i:self.c2_e, self.c1_i:self.c1_e] = Cj
            self.NonResonantRates[self.c2_i:self.c2_e, self.c2_i:self.c2_e] = -Rj
            self.NonResonantRates[self.c3_i:self.c3_e, self.c2_i:self.c2_e] = Rk
            self.NonResonantRates[self.c3_i:self.c3_e, self.c3_i:self.c3_e] = -PEk
            self.NonResonantRates[self.c1_i:self.c1_e, self.c3_i:self.c3_e] = PEi
            
            self.non_resonant_matrix_pop_conservative=Check_population_conservation(self.NonResonantRates,"Non Resonant Rates",self.print_population) and self.non_resonant_matrix_pop_conservative
        self.clk.TOC("Non_Resonant_Rates_Evaluation")

    def Resonant_Rates_Evaluation(self,excitation_freq=[[9076,5]]):
        """
        Description:
            Apply to the 'base_matrix' a transition between two levels 'T' (=[M1,b1,M2,b2,B(Tesla)]) at a rate 'w'
            20120911,9am: execution time around 100 ms
        Todo:
            -Upgrade for several irradiated frequencies at the same time!!!
        Input:
            -k: the transitions index in self.nonzero_i_index,self.nonzero_j_index
        Output:
            -None
        Modification of properties:
            -self.ResonantRates
        Dependence:
            -Properties:
                -self.mw_w
                -self.NonResonantRates
                -self.print_population
                -self.resonant_matrix_pop_conservative
            -Methods:
                -None
        """
        self.clk.TIC("Resonant_Rates_Evaluation")
        if self.fortran_compiled_functions:
            o=np.zeros([self.dim_een+self.dim_een+self.dim_n,self.dim_een+self.dim_een+self.dim_n])
            o=fr.nonresmat(o,self.Mbj_el_wise_square,
                           self.dim_een+self.dim_een+self.dim_n,
                           self.dim_een,
                           self.bulk.dim_PRC,
                           self.donor.dim_e,
                           self.dim_n,
                           self.Cp,
                           self.Cap,
                           self.Rp,
                           self.Rap,
                           self.PE)
            f=excitation_freq[0][0]
            w=excitation_freq[0][1]
            self.ResonantRates=fr.resmat(o,self.Dipolar_Transition_Matrix,
                                         f,w,
                                         self.mw_w,
                                         self.eigvalDzRz,
                                         self.dim_een+self.dim_een+self.dim_n,
                                         self.dim_een)
        else:
            #Print information
            self.Non_Resonant_Rates_Evaluation()
            self.ResonantRates=np.copy(self.NonResonantRates)
            for [freq,width] in excitation_freq:
                lorentzian = lorentz(derivative=0)
                    
                for i in range(self.Dipolar_Transition_Matrix[:, 0].size):
                    for j in range(self.Dipolar_Transition_Matrix[0, :].size):
                        center = np.abs(self.eigvalDzRz[i] - self.eigvalDzRz[j])
                        coeffs = self.Dipolar_Transition_Matrix[i, j] ** 2, center, width
                        #coeffs = 1, center, width
                        evalpts = freq
                        dipole_transition_strength = lorentzian.Eval(coeffs, evalpts)
                        #print center,self.Dipolar_Transition_Matrix[i, j] ** 2
                    
                        i_NRR=self.dim_een+i
                        j_NRR=self.dim_een+j
                    
                        self.ResonantRates[i_NRR,i_NRR]+=-self.mw_w*dipole_transition_strength
                        self.ResonantRates[j_NRR,j_NRR]+=-self.mw_w*dipole_transition_strength
                        self.ResonantRates[i_NRR,j_NRR]+=self.mw_w*dipole_transition_strength
                        self.ResonantRates[j_NRR,i_NRR]+=self.mw_w*dipole_transition_strength
                        
                    #if dipole_transition_strength>1e-6:
                    #    print dipole_transition_strength
                    
            self.resonant_matrix_pop_conservative=Check_population_conservation(self.ResonantRates,"Resonant Rates",self.print_population) and self.resonant_matrix_pop_conservative
        self.clk.TOC("Resonant_Rates_Evaluation")
                
    def EigDpRz(self,bfield=np.array([0,0,1])):
        """
        Description:
            Define the Hamiltonian in the DpRz state
        Todo:
            -None
        Input:
            -B: magnetic field
        Output:
            -self.eigvalDpRz: matrix of the Hamiltonian in the case D+ R0
        Modification of properties:
            -self.hamiltonianDpRz
        Dependence:
            -Properties:
                -self.dim_een
                -self.B
            -Class:
                -self.donor
                -self.bulk
        Comments:
            -relaxation supposed to be faster than others processes
        """
        #nuclear Zeeman and PRC
        #we write only the eigenvalues as it is already diagonalized Assuming B=Bz (here, avoid a diagonalization)
        self.hamiltonianDpRz=np.kron(np.kron(self.bulk.PRC_Hamiltonian_Ze(bfield),self.donor.Sd),self.donor.Id)
        self.hamiltonianDpRz+=np.kron(np.kron(self.bulk.Sd,self.donor.D_Hamiltonian_Ze(bfield)),self.donor.Id)
        self.hamiltonianDpRz+=np.kron(np.kron(self.bulk.Sd,self.donor.Sd),self.donor.D_Hamiltonian_Zn(bfield))

    def EigDzRz(self,bfield=np.array([0,0,1])):
        """
        Description:
            Define the Hamiltonian in the DzRz state
        Todo:
            -None
        Input:
            -None
        Output:
            -self.eigvalDpRz: matrix of the Hamiltonian in the case D+ R0
        Modification of properties:
            -self.hamiltonianDpRz
            -self.Mbj
            -self.Mjb
            -self.Mbj_el_wise_square
            -self.Mjb_el_wise_square
        Dependence:
            -Properties:
                -self.dim_een
                -self.B
                -self.coupling_PRC_donor: Allows B(t) or not
            -Class:
                -self.donor
                -self.bulk
        Comments:
            -relaxation supposed to be faster than others processes
        """
        o=self.donor.D_Hamiltonian_Dz(bfield)
        self.hamiltonianDzRz = np.kron(self.bulk.Sd,o)    #Donor term
        self.hamiltonianDzRz += np.kron(self.bulk.PRC_Hamiltonian_Ze(bfield),np.kron(self.donor.Sd, self.donor.Id)) #PRC term
        
        self.eigvalDzRz, self.Mbj = np.linalg.eig(self.hamiltonianDzRz)  #Diagonalization
        self.Mjb = np.linalg.inv(self.Mbj)
        
        self.Mbj_el_wise_square=self.Mbj**2
        #self.Mjb_el_wise_square=self.Mjb**2

    def EigDpRm(self,bfield=np.array([0,0,1])):
        """
        Description:
            Define the Hamiltonian in the DzRz state
        Todo:
            -None
        Input:
            -B: magnetic field
        Output:
            -None
        Modification of properties:
            -self.hamiltonianDpRm
            -self.Mkb
            -self.Mbk
        Dependence:
            -Properties:
                -self.dim_een
                -self.B
            -Class:
                -self.donor
                -self.bulk
        Comments:
            -relaxation supposed to be faster than others processes
        """
        #nuclear Zeeman and PRC
        #we write only the eigenvalues as it is already diagonalized
        self.hamiltonianDpRm=np.kron(np.kron(self.bulk.Sd,self.donor.Sd),self.donor.D_Hamiltonian_Zn(bfield))
        
    def Do_Eig(self,time_step=0):
        """
        Description:
            Perform the eigensystem, fill the transitions list...
            20120911,9am: execution time around 2 ms
        Todo:
            -None
        Input:
            -B: from MRTools.Bfield(B,xyz)
        Output:
            -None
        Modification of properties:
            -self.Dipolar_Transition_Matrix
            -self.n_of_transitions
            -self.dipolar_transitions
            -self.nonzero_i_index
            -self.nonzero_j_index
        Dependence:
            -Properties:
                -self.Mjb
                -self.Mbj
                -self.Dipolar_Transition_Matrix
            -Methods:
                -self._EigDpRz()
                -self._EigDzRz()
                -self._EigDpRm()
            -class:
                -self.donor
                -self.bulk
        """
        self.CT.TIC("Do_Eig")
        bfield=self.field.field_vec(time_step)
        self.EigDpRz(bfield=bfield)
        self.EigDzRz(bfield=bfield)
        self.EigDpRm(bfield=bfield)
        
        DP1=np.kron(np.kron(self.bulk.Sd,self.donor.Sx),self.donor.Id)    #Dipolar transition
        self.Dipolar_Transition_Matrix=np.dot(self.Mjb,np.dot(DP1,self.Mbj))   #Change the basis
        self.n_of_transitions=np.count_nonzero(self.Dipolar_Transition_Matrix)
        self.dipolar_transitions=np.zeros([self.n_of_transitions,4])    # [[B0,li-lj,EPR],...]
        self.nonzero_i_index,self.nonzero_j_index=np.nonzero(self.Dipolar_Transition_Matrix)
        self.CT.TOC("Do_Eig")

    def Compute_conductivity(self,odeint_sol,time_scale):
        """
        Description:
            Compute the conductivity knowing the populations
        Todo:
            -None
        Input:
            -odeint_sol: solution of the system
            -time_scale=t.size
        Output:
            -self.conductivity
        Modification of properties:
            -self.conductivity
        Dependence:
            -Properties:
                -self.c1_i,self.c1_e
                -self.c3_i,self.c3_e
            -class:
                -self.donor
                -self.bulk
        
        """
        self.CT.TIC("Compute_conductivity")
        self.conductivity=np.zeros([time_scale])
        for k in range(time_scale):
            self.conductivity[k]=self.bulk.mue*np.sum(odeint_sol[k,self.c1_i:self.c1_e])
            self.conductivity[k]+=self.bulk.muh*np.sum(odeint_sol[k,self.c3_i:self.c3_e])
        return self.conductivity
        self.CT.TOC("Compute_conductivity")
          
    def Check_list(self):
        """
        Description:
            Print whether the matrixes are population conservative or not, and if the steady state is reached or not
        Todo:
            -None
        Input:
            -None
        Output:
            -None
        Modification of properties:
            -None
        Dependence:
            -Properties:
                -self.non_resonant_matrix_pop_conservative
                -self.resonant_matrix_pop_conservative
                -self.is_steady_state_reached_flag
        """
        self.CT.TIC("Check_list")
        if not self.non_resonant_matrix_pop_conservative:
            print "\nThe non resonant matrix does not conserve the population."
        
        if not self.resonant_matrix_pop_conservative:
            print "\nAt least one of the resonant matrix does not conserve the population."
        
        if not self.is_steady_state_reached_flag:
            print "\nOne of the steady states at least was not reached."
        self.CT.TOC("Check_list")
  
    def Plot_populations(self,ploted_variable="steady state population",plot_size="full"):
        """
        Input:
            - ploted_variable=      "steady state populations"
                                    "resoning state populations"
                                    "difference populations"
            - plot_size =   "full"
                            "donor"
        """
        self.CT.TIC("Plot_populations")
        if plot_size=="full":
            ind1 = range(self.full_dim)  # the x locations for the groups
        elif plot_size=="donor":
            ind1 = range(self.c2_i,self.c2_e)  # the x locations for the groups
            
        w = 0.4       # the width of the bars
        y1=[]
        for k in ind1:
            if ploted_variable =="steady state populations":
                y1.append(self.steadystate[k])
                title="Steady state populations"
            elif ploted_variable =="resoning state populations":
                y1.append(self.resoningstate[k])
                title="Populations at resonance"
            elif ploted_variable =="difference populations":
                y1.append(self.steadystate[k]-self.resoningstate[k])
                title="Populations at SS - at resonance"
        
        fig=plt.figure()
        ax = fig.add_subplot(1,1,1)
        
        ax.bar(ind1, y1, width=w,color='r')
        
        # add some
        plt.ylabel('Populations')
        plt.title(title)
        
        plt.show()
        self.CT.TOC("Plot_populations")


if __name__ == "__main__":
    #Instanciation
    clk=Computation_time(True)
    clk.TIC("Total")
    sdr_param=SDR_parameters()
    Bi = Shallow_Donor("Bi",clk=clk)
    P = Shallow_Donor("P",clk=clk)
    Si = Silicon(clk=clk)
    sdr_model=Parametrized_SDR_model(Si, Bi, clk=clk)
    transient=Transient(sdr_param,clk=clk)
    
#    SDR.timescale_without_resonance=1e-3
#    SDR.timescale_with_resonance=1e-2
#    SDR.time_intervale_without_resonance=1e-3*SDR.timescale_without_resonance
#    SDR.time_intervale_with_resonance=1e-3*SDR.timescale_with_resonance
    transient.Temporal_evolution_without_resonance(sdr_model)
    transient.Temporal_evolution_with_resonance(sdr_model,excitation_freq=[[9076,1]]) #9076
    transient.Plot_transient(transient_type="field, resonance")# conductivity,light and resonance transient
    
    #SDR.F_sweep_EPR()
    #SDR.coupling_PRC_donor=False
    #SDR.EigDzRz(0)
    
    #SDR.Plot_populations(ploted_variable="difference populations",plot_size="full")#"donor","full"

    toc=time.clock()
    clk.TOC("Total")
    clk.Display()